#include "ECS/System/PBaseMaterialSystem.h"
#include "ECS/PScene.h"
#include "ECS/Component/PLookAtCameraComponent.h"

#include "Graphics/PiccoloVKPipeline.h"
#include "Graphics/PiccoloVKFrameBuffer.h"

#include "Render/PiccoloRenderTarget.h"

#include "PiccoloFileUtil.h"
#include "PiccoloGeometryUtil.h"

namespace Piccolo
{
    void PBaseMaterialSystem::OnInit(PiccoloVKRenderPass *renderPass)
    {
        PiccoloVKDevice* device = GetDevice();
        ShaderLayout shaderLayout =
        {
            .pushConstantRanges =
            {
                {
                    .stageFlags = VK_SHADER_STAGE_VERTEX_BIT,
                    .offset = 0,
                    .size = sizeof(PushConstant)
                }
            }
        };
        mPipelineLayout = std::make_shared<PiccoloVKPipelineLayout>(device, PICCOLO_RES_SHADER_DIR"ecs.vert",
            PICCOLO_RES_SHADER_DIR"ecs.frag", shaderLayout);

        std::vector<VkVertexInputBindingDescription> vertexBindings =
        {
            {
                .binding = 0,
                .stride = sizeof(PiccoloVertex),
                .inputRate = VK_VERTEX_INPUT_RATE_VERTEX
            }
        };
        std::vector<VkVertexInputAttributeDescription> vertexAttributes =
        {
            {
                .location = 0,
                .binding = 0,
                .format = VK_FORMAT_R32G32B32_SFLOAT,
                .offset = offsetof(PiccoloVertex, position)
            },
            {
                .location = 1,
                .binding = 0,
                .format = VK_FORMAT_R32G32_SFLOAT,
                .offset = offsetof(PiccoloVertex, texcoord0)
            },
            {
                .location = 2,
                .binding = 0,
                .format = VK_FORMAT_R32G32B32_SFLOAT,
                .offset = offsetof(PiccoloVertex, normal)
            }
        };

        mPipeline = std::make_shared<PiccoloVKPipeline>(device, renderPass, mPipelineLayout.get());
        mPipeline->SetVertexInputState(vertexBindings, vertexAttributes);
        mPipeline->SetInputAssemblyState(VK_PRIMITIVE_TOPOLOGY_TRIANGLE_LIST)->EnableDepthTest();
        mPipeline->SetDynamicState({ VK_DYNAMIC_STATE_VIEWPORT, VK_DYNAMIC_STATE_SCISSOR });
        mPipeline->SetMultisampleState(VK_SAMPLE_COUNT_4_BIT, VK_FALSE);
        mPipeline->Create();
    }

    void PBaseMaterialSystem::OnRender(VkCommandBuffer cmdBuffer, PiccoloRenderTarget *renderTarget)
    {
        // entt:each
        PScene* scene = GetScene();
        if (!scene) return;

        entt::registry& registry = scene->GetEcsRegistry();
        // 必须要有PTransformComponent PMeshComponent PBaseMaterialComponent这3个基础材质组件
        auto view = registry.view<PTransformComponent, PBaseMaterialComponent>();
        if (view.begin() == view.end()) return;

        // bind pipeline
        mPipeline->Bind(cmdBuffer);

        // setup global params
        PiccoloVKFrameBuffer* frameBuffer = renderTarget->GetCurrentFrameBuffer();
        VkViewport viewport =
        {
            .x = 0,
            .y = 0,
            .width = static_cast<float>(frameBuffer->GetWidth()),
            .height = static_cast<float>(frameBuffer->GetHeight()),
            .minDepth = 0,
            .maxDepth = 1
        };
        vkCmdSetViewport(cmdBuffer, 0, 1, &viewport);
        VkRect2D scissor = { {0, 0}, {frameBuffer->GetWidth(), frameBuffer->GetHeight()} };
        vkCmdSetScissor(cmdBuffer, 0, 1, &scissor);

        glm::mat4 projMatrix = GetProjMatrix(renderTarget);
        glm::mat4 viewMatrix = GetViewMatrix(renderTarget);

        // setup custom params
        view.each([this, &cmdBuffer, &projMatrix, &viewMatrix](const auto& e, const PTransformComponent& transformComponent, const PBaseMaterialComponent& materialComponent)
        {
            auto meshMaterials = materialComponent.GetMeshMaterials();
            for (const auto& entry : meshMaterials)
            {
                PBaseMaterial* material = entry.first;
                if (!material)
                {
                    LOG_W("TODO: default material or error material ? ");
                    continue;
                }
                PushConstant pushConstant
                {
                    .matrix = projMatrix * viewMatrix * transformComponent.GetTransformMatrix(),
                    .colorType = static_cast<uint32_t>(material->colorType)
                };
                vkCmdPushConstants(cmdBuffer, mPipelineLayout->GetPipelineLayout(), VK_SHADER_STAGE_VERTEX_BIT, 0, sizeof(PushConstant), &pushConstant);

                // mesh list draw
                for (const auto& meshIndex : entry.second)
                {
                    PiccoloMesh* mesh = materialComponent.GetMesh(meshIndex);
                    if (mesh)
                    {
                        mesh->Draw(cmdBuffer);
                    }
                }
            }
        });
    }

    void PBaseMaterialSystem::OnDestroy()
    {
        mPipeline.reset();
        mPipelineLayout.reset();
    }
}
